#!/usr/bin/env bash
#! Encoding UTF-8
PATH=/bin:/sbin:/usr/bin:/usr/sbin:/usr/local/bin:/usr/local/sbin:~/bin
export PATH
clear;

trap "echo byebye" EXIT

## 检查中文环境
CN="false"
case $LANG in
  zh_CN*) CN="true";;
esac

[[ -z ${LOG_PATH} ]] \
  && LOG_PATH="${HOME}/log"
[[ ! -d ${LOG_PATH} ]] \
  && mkdir -p ${LOG_PATH}
[[ -z ${INFO_LOG} ]] \
  && INFO_LOG="${LOG_PATH}/tchub_$(date +%Y_%m).log"

if [[ "$(type -t tmp_msg)" != function ]] ; then
  tmp_msg(){
    if [[ -z $2 ]] ; then
      TMP_MSG_RESULT="$1"
    else 
      ${CN} && TMP_MSG_RESULT="$1" || TMP_MSG_RESULT="$2"
    fi
  }
fi

if [[ "$(type -t info)" != function ]] ; then
  info(){
    if [[ -z $2 ]] ; then
      TMP_MSG="$1"
    else 
      ${CN} && TMP_MSG="$1" || TMP_MSG="$2"
    fi
    echo "$(date +'%F %H:%M:%S') ${TMP_MSG} "
  }
fi

if [[ "$(type -t info_log)" != function ]] ; then
  info_log(){
    info "$@" \
      | tee -a ${INFO_LOG}
  }
fi



#====
select_interface(){
  TMP_DEV=""
  INTERFACE_LISTS=($(ip link|awk -F':' '!/^ |lo/ {print $2}'|sed 's/^ //g'))
  select VAR in ${INTERFACE_LISTS[@]} ;do
    TMP_DEV=${VAR}
    [[ -n ${TMP_DEV} ]] \
      && break
  done
}

tc_control_rule_init(){
  DEV=$1
  [[ -n "$2" ]] \
    && RATE_MAX=$2 \
    || RATE_MAX=100
   tc qdisc del dev ${DEV} root
  info_log "rules run \"tc qdisc del dev ${DEV} root\""
  tc qdisc add dev ${DEV} root handle 1: htb default 10
  info_log "rules run \"tc qdisc add dev ${DEV} root handle 1: htb default 10\""
  tc class add dev ${DEV} parent 1: classid 1:1 htb rate ${RATE_MAX}mbit burst 15k
  info_log "rules run \"tc class add dev ${DEV} parent 1: classid 1:1 htb rate ${RATE_MAX}mbit burst 15k\""
}

tc_control_rule(){
  DEV=$1
  CLASS_ID=$2
  RATE_NUM=$3
  CEIL_NUM=$4
  tc class add dev ${DEV} parent 1:1 classid 1:${CLASS_ID} htb rate ${RATE_NUM}mbit ceil ${CEIL_NUM}mbit burst 15k
  info_log "rules run \"tc class add dev ${DEV} parent 1:1 classid 1:${CLASS_ID} htb rate ${RATE_NUM}mbit ceil ${CEIL_NUM}mbit burst 15k\""
  tc qdisc add dev ${DEV} parent 1:${CLASS_ID} handle ${CLASS_ID}: sfq perturb 10
  info_log "rules run \"tc qdisc add dev ${DEV} parent 1:${CLASS_ID} handle ${CLASS_ID}: sfq perturb 10\""
#
}

tc_control_filter_for_ip(){
  DEV=$1
  IP_MATCH=$2
  FLOW_ID=$3
  tc filter add dev ${DEV} protocol ip parent 1:0 prio 1 u32 match ip dst ${IP_MATCH} flowid 1:${FLOW_ID}
  info_log "rules run \"tc filter add dev ${DEV} protocol ip parent 1:0 prio 1 u32 match ip dst ${IP_MATCH} flowid 1:${FLOW_ID}\""
}

tc_control_filter_for_ports(){
  DEV=$1
  PORT_MATCH=$2
  FLOW_ID=$3
  tc filter add dev ${DEV} protocol ip parent 1:0 prio 1 u32 match ip sport ${PORT_MATCH} 0xffff flowid 1:${FLOW_ID}
  info_log "rules run \"tc filter add dev ${DEV} protocol ip parent 1:0 prio 1 u32 match ip sport ${PORT_MATCH} 0xffff flowid 1:${FLOW_ID}\""
}

show_tc_status(){
  DEV=$1
  QDISC_SHOW=$(tc -s -d qdisc show dev ${DEV})
  CLASS_SHOW=$(tc -s -d class show dev ${DEV})
  FILTER_SHOW=$(tc -s -d filter show dev ${DEV})
cat <<EOF
------- qdisc  ---------------
${QDISC_SHOW:-None}

------- class  ---------------
${CLASS_SHOW:-None}

------- filter ---------------
${FILTER_SHOW:-None}

==============================
EOF
}

class_id_max(){
  DEV=$1
  TMP_CLASS_ID=$(tc class show dev ${DEV} |awk '/class htb/ {split($3,attr,":");print attr[2]}'|sort -nr|head -n1)
  MAX_CLASS_ID=$(expr ${TMP_CLASS_ID} + 1)
}

check_num(){
  TMP_NUM=$1
  TMP_BOOL=$(echo ${TMP_NUM}|awk '{start=match($1,/^[0-9]+$/);print start}')
  [[ ${TMP_BOOL} == 0 ]] \
    && info "Input is err , please input agent . " \
    && continue
}

check_ports(){
  TMP_TC_PORTS=$1
  TMP_BOOL=$(echo ${TMP_TC_PORTS}|awk '{start=match($1,/^[1-9][0-9]+$/);print start}')
  if [[ ${TMP_BOOL} == 1 ]];then
    if [[ ${TMP_TC_PORTS} -ge 65535 ]];then
      info "Input is err , please input agent . " \
        && continue
    fi
  else
    info "Input is err , please input agent . " \
      && continue
  fi
}

check_ipaddr(){
  TMP_NUM=$1
  TMP_TC_IPADDR=$(echo ${TMP_NUM} | awk '{start=match($1,/^[1-9][0-9]{0,2}\.[1-9][0-9]{0,2}\.[1-9][0-9]{0,2}\.[1-9][0-9]{0,2}$/);print start}')
  [[ ${TMP_BOOL} == 0 ]] \
    && info "Input is err , please input agent . " \
    && continue
}

select_tc_function(){
  echo "----------------------------------------------------------------"
  declare -a VAR_FUN_LISTS
  info "\033[0;31;1m[Notice]\033[0m 请选择对应的网卡:" "\033[0;31;1m[Notice]\033[0m Select one Interfaces:"
  select_interface
  info "\033[0;31;1m${TMP_DEV} is choose .\033[0m "
  info "\033[0;31;1m[Notice]\033[0m 请选择需要的功能:" "\033[0;31;1m[Notice]\033[0m How to set up iptables:"
  if ${CN} ;then
    VAR_FUN_LISTS=( "退出" "查看${TMP_DEV}的TC状态" "清理${TMP_DEV}的全部TC规则" "生成${TMP_DEV}的TC新规则")
  else
    VAR_FUN_LISTS=("back" "show_tc_status_From_${TMP_DEV}" "Clear_TC_Rules_From_${TMP_DEV}" "Create_New_TC_Rules_From_${TMP_DEV}")
  fi
  select VAR in ${VAR_FUN_LISTS[@]} ;do
    case ${VAR} in
      ${VAR_FUN_LISTS[1]})
        show_tc_status ${TMP_DEV}
        ;;
      ${VAR_FUN_LISTS[2]})
        TMP_OK_RESULT="no"
        tmp_msg \
          "确定清理网卡\"${TMP_DEV}\"上的TC规则吗？(yes or no) [no] : " \
          "Are you sure to clear tc rules from \"${TMP_DEV}\" ？(yes or no) [no] : "
        read -p "${TMP_MSG_RESULT}" -t 30 TMP_OK_RESULT 
        if [[ ${TMP_OK_RESULT} == 'yes' ]];then
          tc qdisc del dev ${DEV} root
          info_log "The TC rules wich ${DEV} is clear !"
        else
          info_log "Nothing to do with ${DEV} !" 
        fi
        ;;
      ${VAR_FUN_LISTS[3]})
        tmp_msg \
          "请设置网卡\"${TMP_DEV}\"上的默认带宽 [100mbit] : " \
          "Set the max rate with \"${TMP_DEV}\" [100mbit] : "
        read -p "${TMP_MSG_RESULT}" -t 30 TC_MAX_RATE 
        [[ -z ${TC_MAX_RATE} ]] \
          && TC_MAX_RATE=100
        check_num ${TC_MAX_RATE}
        tc_control_rule_init ${TMP_DEV} ${TC_MAX_RATE}

        info \
          "请设置网卡\"${TMP_DEV}\"上的规则 : " \
          "Set the rules with \"${TMP_DEV}\" : "

        while true; do
          class_id_max ${TMP_DEV}
          tmp_msg \
            "请设置网卡\"${TMP_DEV}\"上的规则${MAX_CLASS_ID}最低带宽(mbit) : " \
            "Set the min rate in the rules ${MAX_CLASS_ID} with \"${TMP_DEV}\" (mbit) : "
          read -p "${TMP_MSG_RESULT}" TC_MIN_RATE 
          check_num ${TC_MIN_RATE}
          tmp_msg \
            "请设置网卡\"${TMP_DEV}\"上的规则${MAX_CLASS_ID}带宽上限(mbit) : " \
            "Set the max ceil in the rules ${MAX_CLASS_ID} with \"${TMP_DEV}\" (mbit) : "
          read -p "${TMP_MSG_RESULT}" TC_MAX_CEIL
          check_num ${TC_MAX_CEIL}
          tc_control_rule ${TMP_DEV} ${MAX_CLASS_ID} ${TC_MIN_RATE} ${TC_MAX_CEIL}
          tmp_msg \
            "继续添加吗？[no] : " \
            "Continue to add ? [no] : "
          read -p "${TMP_MSG_RESULT}" -t 30 TMP_OK_RESULT 
          [[ -z ${TMP_OK_RESULT} ]] \
            && TMP_OK_RESULT="no"
          [[ ${TMP_OK_RESULT} == 'no' ]] \
            && break
        done

        info \
          "请设置网卡\"${TMP_DEV}\"上的过滤规则 : " \
          "Set the filter rules with \"${TMP_DEV}\" : "

        while true; do
          TMP_FILTER=""
          info \
            "请问针对IP设置还是端口 : " \
            "Set the ip or port : "

          select VAR in "ip" "port" ;do
            TMP_FILTER=${VAR}
            [[ -n ${TMP_FILTER} ]] \
              && break
          done

          CHOOSE_CLASS_ID=""
          echo "=============== class ================"
          tc class show dev ${TMP_DEV} \
            | awk '/class htb/ {split($3,attr,":");print attr[2],"\""$0"\""}'
          echo "======================================"

          info \
            "请问设置在哪个规则上 : " \
            "Which class id would you choose : "

          select VAR in $(tc class show dev ${TMP_DEV} |awk '/class htb/ {split($3,attr,":");print attr[2]}') ;do
            CHOOSE_CLASS_ID=${VAR}
            [[ -n ${CHOOSE_CLASS_ID} ]] \
              && break
          done

          if [[ ${TMP_FILTER} == 'ip' ]];then
            tmp_msg \
              "请输入IP : " \
              "Input the ip : "
            read -p "${TMP_MSG_RESULT}" TMP_TC_IPADDR
            check_ipaddr ${TMP_TC_IPADDR}
            tc_control_filter_for_ip ${TMP_DEV} ${TMP_TC_IPADDR} ${CHOOSE_CLASS_ID}
          else
            tmp_msg \
              "请输入端口号 : " \
              "Input the port num : "
            read -p "${TMP_MSG_RESULT}" TMP_TC_PORT 
            check_ports ${TMP_TC_PORT}
            tc_control_filter_for_ports ${TMP_DEV} ${TMP_TC_PORT} ${CHOOSE_CLASS_ID}
          fi
          tmp_msg \
            "继续添加吗？[no] : " \
            "Continue to add ? [no] : "
          read -p "${TMP_MSG_RESULT}" -t 30 TMP_OK_RESULT 
          [[ -z ${TMP_OK_RESULT} ]] \
            && TMP_OK_RESULT="no"
          [[ ${TMP_OK_RESULT} == 'no' ]] \
            && break
        done
        ;;
      *)
        break
        ;;
    esac
    select_tc_function
  done
}

select_tc_function